import json

import pymysql
import torch
from sklearn.preprocessing import MinMaxScaler

from CnnStockConfig import CnnStockConfig
import numpy as np


class CnnStockDataProvider:

    def __init__(self):
        self.mysql_host = "127.0.0.1"
        self.mysql_port = 3306

        self.connect = pymysql.connect(
            host=self.mysql_host,
            port=self.mysql_port,
            user='root',
            passwd='hxsoft.net',
            database='stock')
        self.cursor = self.connect.cursor()

    def get_stock_data(self):
        """
        获取原始数据
        :return:    原始数据Tensor
        """
        try:
            self.cursor.execute("select * from stock_result where stock_high != ''")
            rows = self.cursor.fetchall()
            if rows is not None:
                seq_list = []
                for row in rows:
                    json_list = json.loads(row[6])
                    item_list = []
                    for index in range(len(json_list)):
                        item = json_list[index]
                        # item_list.append(item['open'])
                        # item_list.append(item['high'])
                        # item_list.append(item['low'])
                        item_list.append(item['close'])
                        item_list.append(item['zl1'])
                        item_list.append(item['zl2'])
                    increment = (100 * (row[3] - row[2])) / row[2]
                    if 0 < increment < 3:
                        item_list.append(1)
                    elif 3 <= increment < 6:
                        item_list.append(3)
                    elif 6 <= increment < 9:
                        item_list.append(5)
                    elif 9 <= increment:
                        item_list.append(7)
                    elif increment == 0:
                        item_list.append(0)
                    elif -3 < increment < 0:
                        item_list.append(2)
                    elif -6 < increment <= -3:
                        item_list.append(4)
                    elif -9 < increment <= -6:
                        item_list.append(6)
                    elif increment <= -9:
                        item_list.append(8)
                    seq_list.append(item_list)
                return seq_list
            else:
                return None
        finally:
            if self.cursor is not None:
                self.cursor.close()
            if self.connect is not None:
                self.connect.close()

    def get_predicated_data(self):
        """
        数据预处理
        :return:
        """
        sep_array = np.array(self.get_stock_data())
        train_size = int(len(sep_array) * CnnStockConfig.ratio)

        train_data = torch.tensor(sep_array[:train_size, 0:len(sep_array[0]) - 1], dtype=torch.float32)
        train_value = torch.tensor(sep_array[:train_size, len(sep_array[0]) - 1], dtype=torch.long)

        test_data = torch.tensor(sep_array[train_size:, 0:len(sep_array[0]) - 1], dtype=torch.float32)
        test_value = torch.tensor(sep_array[train_size:, len(sep_array[0]) - 1], dtype=torch.long)

        scaler1 = MinMaxScaler(feature_range=(0, 1))
        train_data = torch.tensor(scaler1.fit_transform(train_data), dtype=torch.float32)
        test_data = torch.tensor(scaler1.fit_transform(test_data), dtype=torch.float32)

        return train_data, train_value, test_data, test_value
